#a simple network-based SIR model written in Python
#Jon Zelner
#University of Michigan
#August 11, 2009


import igraph
import random
import copy
import pylab as pl

class simpleNetworkSIRModel():
    def __init__(self, b = .2,  g = .01, N = 300, p = .02):
        self.b = b
        self.g = g
        self.t = 0

        self.N = N
        self.graph = igraph.Graph.Watts_Strogatz(1, N, nei=8, p = .02)
        self.graph.simplify()

        self.adjacencyList = []
        self.sAgentList = []
        self.iAgentList = []
        self.rAgentList = []
        
        self.sList = []
        self.iList = []
        self.rList = []
        self.newIList = []

        for i in range(self.N):
            self.adjacencyList.append([])
            
        for edge in self.graph.es:
            self.adjacencyList[edge.source].append(edge.target)
            self.adjacencyList[edge.target].append(edge.source)
        
        
        allAgents = range(N)
        random.shuffle(allAgents)

        self.sAgentList = copy.copy(allAgents)
        self.indexCase = self.sAgentList.pop()
        self.iAgentList.append(self.indexCase)

    def run(self):
        while len(self.iAgentList) > 0:
            tempIAgentList = []
            recoverList = []
            newI = 0
            for iAgent in self.iAgentList:
                for agent in self.adjacencyList[iAgent]:
                    if agent in self.sAgentList:
                        if (random.random() < self.b):
                            self.sAgentList.remove(agent)
                            tempIAgentList.append(agent)
                            newI += 1
                if (random.random() < self.g):
                    recoverList.append(iAgent)
            
            for recoverAgent in recoverList:
                self.iAgentList.remove(recoverAgent)
                self.rAgentList.append(recoverAgent)
                
            self.iAgentList.extend(tempIAgentList)
            self.sList.append(len(self.sAgentList))
            self.iList.append(len(self.iAgentList))
            self.rList.append(len(self.sAgentList))
            self.newIList.append(newI)
            self.t += 1
            random.shuffle(self.iAgentList)

        return [self.sList, self.iList, self.rList, self.newIList]


    def graphPlot(self):
        for v in self.graph.vs():
            v['label_size'] = 0
            v['color'] = 'blue'
            if v.index in self.rAgentList or v.index in self.iAgentList:
                v['color'] = 'red'
            if v.index == self.indexCase:
                v['color'] = 'green'
        
        l = self.graph.layout_kamada_kawai()
        igraph.drawing.plot(self.graph, layout = l)
    
    def iPlot(self):
        pl.figure()
        pl.plot(self.iList)
        pl.show()

    
            